import numpy as np
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
from sklearn.datasets import make_blobs
import random


class KMeans:
    def __init__(self, k=3, max_iters=100, tol=1e-4, random_state=None):
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
        self.random_state = random_state
        self.centroids = None
        self.labels = None
        self.inertia_ = None

    def _initialize_centroids(self, X):
        """初始化质心 - 使用K-means++方法"""
        n_samples, n_features = X.shape

        if self.random_state is not None:
            np.random.seed(self.random_state)

        # 第一个质心随机选择
        centroids = [X[np.random.randint(n_samples)]]

        # 选择剩余的k-1个质心
        for _ in range(1, self.k):
            # 计算每个样本到最近质心的距离
            distances = np.array([min([np.linalg.norm(x - c) for c in centroids])
                                  for x in X])

            # 根据距离的概率分布选择下一个质心
            probabilities = distances ** 2 / (distances ** 2).sum()
            cumulative_probs = probabilities.cumsum()
            r = np.random.rand()

            for i, p in enumerate(cumulative_probs):
                if r < p:
                    centroids.append(X[i])
                    break

        return np.array(centroids)

    def _assign_clusters(self, X, centroids):
        """将样本分配到最近的质心"""
        n_samples = X.shape[0]
        labels = np.zeros(n_samples, dtype=int)

        for i in range(n_samples):
            distances = [np.linalg.norm(X[i] - centroid) for centroid in centroids]
            labels[i] = np.argmin(distances)

        return labels

    def _update_centroids(self, X, labels):
        """更新质心位置"""
        centroids = np.zeros((self.k, X.shape[1]))

        for i in range(self.k):
            # 获取属于第i个簇的所有样本
            cluster_points = X[labels == i]
            if len(cluster_points) > 0:
                centroids[i] = cluster_points.mean(axis=0)
            else:
                # 如果簇为空，随机重新初始化
                centroids[i] = X[np.random.randint(X.shape[0])]

        return centroids

    def _compute_inertia(self, X, labels, centroids):
        """计算簇内平方和（ inertia ）"""
        inertia = 0
        for i in range(self.k):
            cluster_points = X[labels == i]
            if len(cluster_points) > 0:
                inertia += np.sum((cluster_points - centroids[i]) ** 2)
        return inertia

    def fit(self, X):
        n_samples, n_features = X.shape

        # 初始化质心
        self.centroids = self._initialize_centroids(X)

        # 迭代优化
        for iteration in range(self.max_iters):
            # 分配样本到簇
            labels = self._assign_clusters(X, self.centroids)

            # 更新质心
            new_centroids = self._update_centroids(X, labels)

            # 计算质心移动距离
            centroid_shift = np.linalg.norm(new_centroids - self.centroids, axis=1).max()

            # 更新质心
            self.centroids = new_centroids
            self.labels = labels

            # 计算inertia
            self.inertia_ = self._compute_inertia(X, labels, self.centroids)

            # 检查收敛
            if centroid_shift < self.tol:
                print(f"收敛于第 {iteration + 1} 次迭代")
                break

        return self

    def predict(self, X):
        """预测新样本的簇标签"""
        return self._assign_clusters(X, self.centroids)

    def fit_predict(self, X):
        """训练模型并返回簇标签"""
        self.fit(X)
        return self.labels


# 测试和可视化
def test_kmeans():
    # 生成测试数据
    X, y_true = make_blobs(n_samples=300, centers=3, n_features=2,
                           random_state=42, cluster_std=0.60)

    # 创建KMeans实例
    kmeans = KMeans(k=3, max_iters=100, random_state=42)

    # 训练模型
    labels = kmeans.fit_predict(X)

    # 可视化结果
    plt.figure(figsize=(12, 5))

    # 原始数据
    plt.subplot(1, 2, 1)
    plt.scatter(X[:, 0], X[:, 1], c=y_true, cmap='viridis', s=50, alpha=0.8)
    plt.title('原始数据 (真实标签)')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')

    # 聚类结果
    plt.subplot(1, 2, 2)
    plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis', s=50, alpha=0.8)
    plt.scatter(kmeans.centroids[:, 0], kmeans.centroids[:, 1],
                c='red', marker='X', s=200, label='质心')
    plt.title('KMeans 聚类结果')
    plt.xlabel('特征 1')
    plt.ylabel('特征 2')
    plt.legend()

    plt.tight_layout()
    plt.show()

    print(f"簇内平方和 (Inertia): {kmeans.inertia_:.2f}")
    print("质心位置:")
    for i, centroid in enumerate(kmeans.centroids):
        print(f"簇 {i}: {centroid}")


# 肘部法则确定最佳K值
def elbow_method(X, max_k=10):
    """使用肘部法则确定最佳聚类数量"""
    inertias = []
    k_range = range(1, max_k + 1)

    for k in k_range:
        kmeans = KMeans(k=k, random_state=42)
        kmeans.fit(X)
        inertias.append(kmeans.inertia_)

    # 绘制肘部图
    plt.figure(figsize=(8, 6))
    plt.plot(k_range, inertias, 'bo-')
    plt.xlabel('聚类数量 K')
    plt.ylabel('簇内平方和 (Inertia)')
    plt.title('肘部法则 - 确定最佳K值')
    plt.grid(True)
    plt.show()

    return inertias


if __name__ == "__main__":
    # 生成测试数据
    X, _ = make_blobs(n_samples=300, centers=3, n_features=2,
                      random_state=42, cluster_std=0.60)

    # 测试KMeans
    test_kmeans()

    # 使用肘部法则
    print("\n使用肘部法则确定最佳K值:")
    inertias = elbow_method(X)